import os
from langchain_community.document_loaders import WebBaseLoader
from langchain_text_splitters import RecursiveCharacterTextSplitter
from langchain_core.vectorstores import InMemoryVectorStore
from langchain_community.embeddings import HuggingFaceEmbeddings  # 改用本地模型
from langchain_community.chat_models import ChatZhipuAI
from langchain_core.prompts import ChatPromptTemplate
from langchain_core.runnables import RunnablePassthrough
from langchain_core.output_parsers import StrOutputParser
from dotenv import load_dotenv

# 加载环境变量
load_dotenv()

model_name = "D:/ideaSpace/MyPython/models/bge-small-zh-v1.5"

class StableURLRAGSystem:
    def __init__(self, url: str):
        """
        稳定版URL RAG系统

        参数:
            url: 要处理的网页URL
        """
        self.url = url

        # 设置USER_AGENT防止抓取警告
        os.environ["USER_AGENT"] = "Mozilla/5.0 (Windows NT 10.0; Win64; x64) AppleWebKit/537.36"

        # 初始化组件（改用本地嵌入模型）
        self.embeddings = HuggingFaceEmbeddings(
            model_name=model_name,
            model_kwargs={'device': 'cpu'}
        )

        self.llm = ChatZhipuAI(
            api_key=os.getenv("ZHIPUAI_API_KEY"),
            model="glm-4",
            temperature=0.3
        )

        # 执行处理流程
        self._process_pipeline()

    def _process_pipeline(self):
        """完整的处理流程"""
        try:
            # 1. 加载网页内容
            documents = self._load_web_content()

            # 2. 创建内存向量存储
            self.vectorstore = InMemoryVectorStore.from_documents(
                documents=documents,
                embedding=self.embeddings
            )

            # 3. 构建RAG链
            self.rag_chain = (
                    {"context": self.vectorstore.as_retriever(), "question": RunnablePassthrough()}
                    | ChatPromptTemplate.from_template("基于以下内容回答:{context}\n问题:{question}")
                    | self.llm
                    | StrOutputParser()
            )

        except Exception as e:
            raise RuntimeError(f"处理失败: {str(e)}")

    def _load_web_content(self):
        """加载并分割网页内容"""
        loader = WebBaseLoader(
            self.url,
            requests_kwargs={"headers": {"User-Agent": os.getenv("USER_AGENT")}}
        )

        text_splitter = RecursiveCharacterTextSplitter(
            chunk_size=800,
            chunk_overlap=100,
            separators=["\n\n", "\n", "。", "！", "？"]
        )

        return text_splitter.split_documents(loader.load())

    def query(self, question: str) -> str:
        """执行查询"""
        try:
            return self.rag_chain.invoke(question)
        except Exception as e:
            return f"查询出错: {str(e)}"

# 使用示例
if __name__ == "__main__":
    try:
        # 初始化系统（替换为实际URL）
        rag = StableURLRAGSystem(url="https://blog.csdn.net/weixin_45690643/article/details/136328902")

        # 交互式查询
        while True:
            question = input("\n请输入问题(输入q退出): ").strip()
            if question.lower() == 'q':
                break
            print(f"\n回答：{rag.query(question)}")

    except Exception as e:
        print(f"系统错误: {str(e)}")